"""
a local server where images are not saved.
"""
import os
# from turtle import distance
import uvicorn
import base64
import numpy as np
import cv2, time
import hashlib
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from engine import SearchEngine, UploadParams, SearchParams, file_to_base64
from engine.utils import get_config


ERROR_CODE_MAP = {
    4001: '请求参数错误',
    4002: '图像解析失败'
}


def get_error_message(code):
    if code == 0:
        return {'code': 0, 'message': 'success'}
    message = ERROR_CODE_MAP[code]
    return {'code': code, 'message': message}


def data_md5(data):
    md5_obj = hashlib.md5()
    md5_obj.update(data)
    code = md5_obj.hexdigest()
    return str(code).lower()


app = FastAPI()
origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# config_path = 'cfg/mocov3.yaml'
config_path = 'cfg/resnet.yaml'
config = get_config(config_path)
if not os.path.exists(config.IMAGE_SAVE_FOLDER):
    os.makedirs(config.IMAGE_SAVE_FOLDER)

engine = SearchEngine(config)


@app.post('/hello')
async def hello():
    return {'code': 0, 'message': 'hello'}


@app.post('/img/upload')
async def upload(data: UploadParams):
    data = dict(data)
    timestamp = data.get('timestamp')
    if not timestamp:
        return get_error_message(4001)
    if 'image' not in data.keys():
        return get_error_message(4001)
    try:
        image_data = data['image']
        image_data = bytes(image_data, encoding='utf-8')
        decoded = base64.urlsafe_b64decode(image_data)
        decoded = np.fromstring(decoded, np.uint8)
        image = cv2.imdecode(decoded, cv2.IMREAD_COLOR)
    except Exception as e:
        return get_error_message(4002)
    
    # save image into local disk drive
    image_path = os.path.join(config.IMAGE_SAVE_FOLDER, '{}.jpg'.format(data_md5(image_data)))
    if not os.path.exists(image_path):
        cv2.imwrite(image_path, image)

    id = engine.upload(config.TABLE_NAME, image_path, timestamp)
    return {'code': 0, 'message': {"id": id}}


@app.post('/img/search')
async def search(data: SearchParams):
    data = dict(data)
    if 'image' not in data.keys():
        return get_error_message(4001)
    
    # decode image data
    try:
        image_data = data['image']
        image_data = bytes(image_data, encoding='utf-8')
        decoded = base64.urlsafe_b64decode(image_data)
        decoded = np.fromstring(decoded, np.uint8)
        image = cv2.imdecode(decoded, cv2.IMREAD_COLOR)
    except Exception as e:
        return get_error_message(4002)

    image_paths, distances, ids = engine.search(config.TABLE_NAME, image, top_k=data['topk'], start_time=data["start_time"],
                                            end_time=data["end_time"])
    if config.MILVUS.METRIC_TYPE == 'L2':
        distances = [np.sqrt(v) for v in distances]
        distances = [1.0 - d for d in distances]
    elif config.MILVUS.METRIC_TYPE == 'IP':
        pass
    else:
        raise ValueError("invalid metric type")

    # construct image urls with nginx
    # image_urls = []
    # base_url = 'http://180.114.8.107:8325/image_search/'
    # for p in result_paths:
    #     filename = os.path.basename(p)
    #     image_urls.append(base_url + filename)

    return {'code': 0, 'message': {'image_data': image_paths, 'distances': distances, 'ids': ids}}


if __name__ == '__main__':
    uvicorn.run(app=app, host='0.0.0.0', port=1088)
